import os
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
from GLOB import PICTURE_SIZE


transform = transforms.Compose([
    transforms.Resize((PICTURE_SIZE, PICTURE_SIZE)),  # 调整图片大小
    transforms.ToTensor(),           # 将图片转换为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = [os.path.join(root_dir, f) for f in os.listdir(root_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image_path = self.images[idx]
        image = Image.open(image_path).convert('RGB')  # 确保图片是RGB格式
        if self.transform:
            image = self.transform(image)
        return image

def show():
    # 将Tensor转换回PIL图像以便显示
        image = image.numpy().transpose((1, 2, 0))  # 转换为 (H, W, C) 格式
        image = image * 0.5 + 0.5  # 反归一化
        
        # 显示图片
        plt.imshow(image)
        plt.title('Sample Image')
        plt.show()

if __name__=='__main__':

    root_dir = r'data\butterfly'  # 图片文件夹路径
    dataset = CustomDataset(root_dir, transform=transform)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

    for images in dataloader:
        # 选择显示第一张图片
        image = images[0]  # 获取第一张图片的Tensor
        
        # 将Tensor转换回PIL图像以便显示
        image = image.numpy().transpose((1, 2, 0))  # 转换为 (H, W, C) 格式
        image = image * 0.5 + 0.5  # 反归一化
        
        # 显示图片
        plt.imshow(image)
        plt.title('Sample Image')
        plt.show()
        
        # 你可以在这里进行模型训练或测试
        break  # 只显示一张图片，然后退出循环